from typing import TypeVar, Generic, List

from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from ollama import Tool
from pydantic import BaseModel

T = TypeVar('T')

class CompanyName(BaseModel):
    name: str


class ValueWrap(BaseModel, Generic[T]):
    value: T

llm = init_chat_model("qwen2:7b",model_provider='ollama',
                api_key="EMPTY",
                api_base="http://localhost:11434",
                )
from langchain_core.tools import tool



template = ("角色:我是物流专家，"
            "任务：生成包含5个 name 的列表,name为公司名称，公司名称必须是中文"
            "要求:"
      "  - 返回数据必须通过tool_calls加工"
            " - 公司名称具有新颖性和独创性"
            "上下文:"
            "   {{content}}"
            )

structured_llm = llm.with_structured_output(schema=ValueWrap[List[str]], include_raw=True,method="function_calling")


prompt = ChatPromptTemplate([
        ("system", template),
        ("human", "输出:"),
    ], template_format="jinja2")

chain = prompt | structured_llm
resp = chain.invoke(input='',content='好听的')
return_data = resp['raw'].tool_calls[0]['args']['value']

print(return_data)

